import pandas as pd
import tushare as ts
from datetime import datetime, timedelta
from typing import Optional, Dict, List
import logging
from trading.core.data_model import MarketData

logger = logging.getLogger(__name__)


class DataFetcher:
    """股票数据获取器（兼容Tushare Pro最新API）"""

    def __init__(self, ts_api: ts.pro_api, data_queue: Optional[object] = None):
        self.ts_api = ts_api
        self.data_queue = data_queue

    def fetch_data(self, stock_code: str) -> Optional[MarketData]:
        """获取指定股票的完整市场数据"""
        try:
            # 1. 确定市场代码
            market = "SH" if stock_code.startswith('6') else "SZ"
            ts_code = f"{stock_code}.{market}"

            # 2. 获取实时行情（新版API）
            quote = self._get_realtime_quote(ts_code)
            if quote.empty:
                logger.warning(f" 未获取到{stock_code}的实时行情")
                return None

                # 3. 获取各类K线数据（使用标准API）
            kline_data = self._get_all_kline_data(ts_code)

            # 4. 构建MarketData对象
            data = MarketData(
                price=float(quote['current_price'].iloc[0]),
                open=float(quote['open'].iloc[0]),
                close=float(quote['pre_close'].iloc[0]),
                high=float(quote['high'].iloc[0]),
                low=float(quote['low'].iloc[0]),
                volume=int(quote['volume'].iloc[0]),
                amount=float(quote['amount'].iloc[0]),
                change=float(quote['change'].iloc[0]),
                p_change=float(quote['pct_chg'].iloc[0]),
                tick_data=self._get_tick_data(ts_code),
                buy_orders=self._parse_orders(quote, 'bid', 5),
                sell_orders=self._parse_orders(quote, 'ask', 5),
                deals=self._get_deals(ts_code),
                kline_data=kline_data
            )
            return data

        except Exception as e:
            logger.error(f" 数据获取失败: {str(e)}", exc_info=True)
            return None

    def _get_realtime_quote(self, ts_code: str) -> pd.DataFrame:
        """使用新版API获取实时行情"""
        try:
            # 标准API调用（需要基础权限）
            return self.ts_api.daily(ts_code=ts_code,
                                     trade_date=datetime.now().strftime('%Y%m%d'))
        except Exception as e:
            logger.warning(f" 实时行情获取失败，使用备用方案: {str(e)}")
            # 备用方案：旧版API
            code = ts_code.split('.')[0]
            return ts.get_realtime_quotes(code)

    def _get_tick_data(self, ts_code: str) -> pd.DataFrame:
        """获取当日分时数据（使用标准API）"""
        try:
            # 标准API调用
            return self.ts_api.minute(ts_code=ts_code,
                                      freq='1min',
                                      trade_date=datetime.now().strftime('%Y%m%d'))
        except Exception as e:
            logger.warning(f" 分时数据获取失败: {str(e)}")
            return pd.DataFrame()

    def _get_all_kline_data(self, ts_code: str) -> Dict[str, pd.DataFrame]:
        """获取多周期K线数据（使用标准API）"""
        kline_data = {}
        end_date = datetime.now().strftime('%Y%m%d')
        start_date = (datetime.now() - timedelta(days=365)).strftime('%Y%m%d')

        # 各周期对应的API映射
        api_mapping = {
            '5min': ('minute', '5min'),
            '15min': ('minute', '15min'),
            '30min': ('minute', '30min'),
            '60min': ('minute', '60min'),
            'day': ('daily', None),
            'week': ('weekly', None),
            'month': ('monthly', None)
        }

        for key, (api_name, freq) in api_mapping.items():
            try:
                # 动态调用API
                api_func = getattr(self.ts_api, api_name)
                params = {'ts_code': ts_code, 'start_date': start_date, 'end_date': end_date}
                if freq: params['freq'] = freq

                data = api_func(**params)

                if not data.empty:
                    data['trade_date'] = pd.to_datetime(data['trade_date'])
                    data.set_index('trade_date', inplace=True)
                    data.rename(columns={
                        'open': 'Open', 'high': 'High', 'low': 'Low',
                        'close': 'Close', 'vol': 'Volume'
                    }, inplace=True)
                    kline_data[key] = data
            except Exception as e:
                logger.warning(f" 获取{key}K线失败: {str(e)}")

        return kline_data

    def _parse_orders(self, quote: pd.DataFrame, side: str, count: int) -> List[Dict]:
        """解析盘口数据（兼容新版字段）"""
        orders = []
        for i in range(1, count + 1):
            # 新版字段格式：bid_price1, bid_vol1 等
            price_key = f"{side}_price{i}"
            vol_key = f"{side}_vol{i}"

            # 旧版兼容：b1_p, b1_v 等
            if price_key not in quote.columns:
                price_key = f"{side}{i}_p"
                vol_key = f"{side}{i}_v"

            orders.append({
                "price": float(quote.get(price_key, [0]).iloc[0]),
                "volume": int(quote.get(vol_key, [0]).iloc[0])
            })
        return orders

    def _get_deals(self, ts_code: str) -> pd.DataFrame:
        """获取成交明细（新版兼容方案）"""
        try:
            # 标准API调用
            return self.ts_api.tick(ts_code=ts_code,
                                    trade_date=datetime.now().strftime('%Y%m%d'))
        except Exception as e:
            logger.warning(f" 成交明细获取失败: {str(e)}")
            # 备用方案
            code = ts_code.split('.')[0]
            return ts.get_today_ticks(code)